﻿using System;
using System.Collections.Generic;
using ABMath.ModelFramework.Data;
using ABMath.ModelFramework.Models;
using ABMath.ModelFramework.Transforms;
using MathNet.Numerics;
using NUnit.Framework;

namespace ABMath.Tests.ModelFramework.Models
{
    [TestFixture]
    public class ARMAModelTest
    {
        private ARMAModel model1;
        private ARMAModel model2;

        private TimeSeries simulatedData;

        [SetUp]
        public void SetUp()
        {
            // create a couple of models
            model1 = new ARMAModel(1, 0);
            model1.SetARPolynomial(new Polynomial(new[] {1.0, -0.5}));
            
            model2 = new ARMAModel(0, 1);
            model2.SetMAPolynomial(new Polynomial(new[] { 1.0, -0.5 }));

            // simulate a time series from one of them
            var times = new List<DateTime>();
            var current = new DateTime(2001, 1, 1);  // Jan. 1st, 2001
            for (int t = 0; t < 1000; ++t)            // create list of 100 days starting with Jan. 1st, 2001
            {
                times.Add(current);
                current = current.AddDays(1);
            }
            simulatedData = model2.SimulateData(times, 0) as TimeSeries; // simulate into times in list, using random number seed 0
        }

        [Test]
        public void YuleWalkerTest()
        {
            var acvf = simulatedData.ComputeACF(4, false);
            ARMAModel ma1 = new ARMAModel(0, 1);
            ma1.EstimateByYuleWalker(0, acvf);
            ARMAModel arma11 = new ARMAModel(1, 1);
            arma11.EstimateByYuleWalker(0, acvf);
        }

        [Test]
        public void FittingAndForecastingTest()
        {
            // fit model first, using maximum likelihood estimation
            var model = new ARMAModel(1, 1);     // create the model object
            model.TheData = simulatedData;       // this is the data we want to fit the model to
            model.FitByMLE(200, 100, 0, null);   // first param is # low discrepancy sequence iterations, should be at least about 200
                                                 // second param is # standard optimizer iterations, should be at least about 100

            // now do some forecasting beyond the end of the data
            var forecaster = new ForecastTransform();

            var futureTimes = new List<DateTime>();
            var nextTime = simulatedData.GetLastTime();
            for (int t = 0; t < 8; ++t )                    // go eight days into the future beyond the end of the data we have
            {
                nextTime = nextTime.AddDays(1);
                futureTimes.Add(nextTime);
            }
            forecaster.FutureTimes = futureTimes.ToArray(); // these are future times at which we want forecasts
                                                            // for now, ARMA models do not use the times in any meaningful way when forecasting,
                                                            // they just assume that these are the timestamps for the next sequential points
                                                            // after the end of the existing time series

            forecaster.SetInput(0, model, null);            // the ARMA model used for forecasting
            forecaster.SetInput(1, simulatedData, null);    // the original data

            // normally you would call the Recompute() method of a transform, but there is no need
            // here; it is automatically called as soon as all inputs are set to valid values (in the 2 statements above)
            var predictors = forecaster.GetOutput(0) as TimeSeries;

            // now predictors is a time series of the forecast values for the next 8 days
            // that is, predictors[0] is the predictive mean of X_{101} given X_1,...,X_100,
            //          predictors[1] is the predictive mean of X_{102} given X_1,...,X_100, etc.
        }

        [Test]
        public void RealTimeTest()
        {
            model1.ResetRealTimePrediction();
            model2.ResetRealTimePrediction();

            var current = new DateTime(2001, 1, 1);

            model1.Register(current, 1.0);
            Assert.AreEqual(model1.GetCurrentPredictor(current).Mean, 0.5);

            model2.Register(current, 1.0);
            double x = model2.GetCurrentPredictor(current).Mean;
            Assert.AreEqual(x, -0.4);

            model1.Register(current, -1.0);
            Assert.AreEqual(model1.GetCurrentPredictor(current).Mean, -0.5);

            model2.Register(current, -1.0);
            x = model2.GetCurrentPredictor(current).Mean;

            model1.Register(current, -2.0);
            Assert.AreEqual(model1.GetCurrentPredictor(current).Mean, -1.0);

            model2.Register(current, -2.0);
            x = model2.GetCurrentPredictor(current).Mean;
        }
    }
}
